# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2025/5/27 13:37
# @Function:
"""
我有两份tif数据
一份是：生态图.tif中包含很多生态类型（森林、草地、湿地）；
一份是：空气净化的核算结果.tif;
现在需要根据生态图的类型将空气净化的核算结果.tif拆分出来；
比如将空气净化的核算结果.tif中只有森林生态系统的保存成一份空气净化的核算结果_森林.tif
"""
import rasterio
import numpy as np
import os

# 定义生态系统类型与生态图编码的映射（来自 formulas.py）
from calc_v1.utils import load_parameters_from_json, extract_codes


eco_type_json_path = r"F:\code\dev\calc-gep-regulate-cqc\calc_v1\data\ecosystems_json\ecosystems_type_GB.json"
# 读取 JSON 数据
ecosystem_json = load_parameters_from_json(eco_type_json_path)

# 为每个一级生态系统生成Code列表
eco_mapping = {}
for ecosystem in ecosystem_json["applicable_ecosystems"]:
    name = ecosystem["name"]
    codes = extract_codes(ecosystem)
    eco_mapping[name] = codes

# 输入文件路径
root_path = r"H:\test\eco_classify"
eco_tif_path = os.path.join(root_path, "china_110000.tif")
air_cleaning_tif_path = os.path.join(root_path, "110000_Qsr_Qsr_replaced_output.tif")

# 输出目录
output_dir = r"H:\test\eco_classify\output"
os.makedirs(output_dir, exist_ok=True)

# 读取生态图数据
with rasterio.open(eco_tif_path) as eco_src:
    eco_data = eco_src.read(1)
    eco_nodata = eco_src.nodata
    eco_profile = eco_src.profile
    eco_transform = eco_src.transform
    eco_crs = eco_src.crs

# 读取空气质量数据
with rasterio.open(air_cleaning_tif_path) as air_src:
    air_data = air_src.read(1)
    air_nodata = air_src.nodata
    air_profile = air_src.profile

# 检查分辨率和范围是否一致
assert eco_data.shape == air_data.shape, "生态图与土壤保持图尺寸不一致！请确保它们的空间对齐。"

# 创建生态图的有效区域掩膜（排除 NoData）
if eco_nodata is not None:
    eco_valid_mask = eco_data != eco_nodata
else:
    eco_valid_mask = np.ones_like(eco_data, dtype=bool)

# 创建空气质量图的有效区域掩膜（排除 NoData）
if air_nodata is not None:
    air_valid_mask = air_data != air_nodata
else:
    air_valid_mask = np.ones_like(air_data, dtype=bool)

# 合并两个掩膜：仅保留生态图和空气质量图都有效的区域
valid_mask = eco_valid_mask & air_valid_mask

# 遍历每个生态系统类型
for eco_name, eco_codes in eco_mapping.items():
    # 创建当前生态系统的掩膜
    eco_mask = np.isin(eco_data, eco_codes)

    # 合并最终掩膜：同时满足是该生态系统、且数据有效
    final_mask = eco_mask & valid_mask

    # 提取空气质量数据，无效区域设为 NaN
    extracted_data = np.where(final_mask, air_data, np.nan)

    if np.isnan(extracted_data).all():
        print(f"{eco_name} 在生态图中无对应区域，跳过保存。")
        continue

    # 构建输出文件路径
    base_filename = os.path.basename(air_cleaning_tif_path)  # 获取原始文件名，如 "110000_Qsr_Qsr_replaced_output.tif"

    # 分割文件名
    parts = base_filename.split("_")  # 按下划线分割
    # 替换第三个字段（索引为2），例如把 Qsr 替换成 森林生态系统
    if len(parts) > 2:
        parts[2] = eco_name  # 替换第3段为生态系统名称
    else:
        raise ValueError("文件名格式不符合预期，无法找到第3个字段")

    # 重新拼接文件名
    new_filename = "_".join(parts)  # 如 "110000_Qsr_森林生态系统_replaced_output.tif"
    output_path = os.path.join(output_dir, new_filename)

    # 更新 profile 以支持写入 NaN 值
    eco_profile.update(
        dtype=rasterio.float32,
        nodata=air_nodata,
        count=1,
        crs=eco_crs,
        transform=eco_transform,
        driver='GTiff',
        compress='lzw',
        bigtiff='YES'
    )

    # 写入新的 TIFF 文件
    with rasterio.open(output_path, 'w', **eco_profile) as dst:
        dst.write(extracted_data.astype(np.float32), 1)

    print(f"已保存 {eco_name} 的土壤保持结果到 {output_path}")
